import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from sklearn.metrics import pairwise_distances
import numpy as np

# 计算CKA相似度
def cka(X, Y):
    # 计算核矩阵
    K_X = X @ X.T
    K_Y = Y @ Y.T
    # 中心化
    ones = torch.ones(K_X.size(0), K_X.size(0), device=K_X.device) / K_X.size(0)
    K_X_centered = K_X - ones @ K_X - K_X @ ones + ones @ K_X @ ones
    K_Y_centered = K_Y - ones @ K_Y - K_Y @ ones + ones @ K_Y @ ones
    # 计算弗罗贝纽斯内积
    frobenius_inner_product = torch.trace(K_X_centered @ K_Y_centered.T)
    # 计算范数
    norm_K_X = torch.sqrt(torch.trace(K_X_centered @ K_X_centered.T))
    norm_K_Y = torch.sqrt(torch.trace(K_Y_centered @ K_Y_centered.T))
    return frobenius_inner_product / (norm_K_X * norm_K_Y)

# 定义模型
model = torchvision.models.resnet18(pretrained=True)
model.fc = nn.Linear(model.fc.in_features, 10)  # 修改最后一层

# 使用GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# 加载数据集
transform = transforms.Compose([
    transforms.Resize(128),
    transforms.ToTensor(),
])

# 假设数据集是 CIFAR-10
train_dataset = torchvision.datasets.CIFAR10(root='/home/data/CIFAR10', train=False, download=False, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)

# 设置优化器
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
criterion = nn.CrossEntropyLoss()

# 预训练网络特征提取函数
def get_features(model, data_loader):
    model.eval()
    features = []
    with torch.no_grad():
        for images, _ in data_loader:
            images = images.to(device)
            x = model.conv1(images)
            x = model.bn1(x)
            x = model.relu(x)
            x = model.maxpool(x)
            x = model.layer1(x)
            x = model.layer2(x)
            x = model.layer3(x)
            x = model.layer4(x)
            x = x.mean([2, 3])  # Global average pooling
            features.append(x.cpu().numpy())
    return np.concatenate(features, axis=0)

# 计算原始预训练模型的特征表示
train_features_before = get_features(model, train_loader)

# 微调模型并在每个epoch输出相似度
num_epochs = 10
for epoch in range(num_epochs):
    model.train()
    for i, (images, labels) in enumerate(train_loader):
        images, labels = images.to(device), labels.to(device)
        
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

    # 计算当前模型的特征表示
    train_features_after = get_features(model, train_loader)

    # 计算CKA相似度
    similarity = cka(torch.tensor(train_features_before), torch.tensor(train_features_after))
    print(f"Epoch {epoch+1}/{num_epochs}, CKA similarity: {similarity.item():.4f}")

    # 更新原始特征
    train_features_before = train_features_after
